import pdfplumber
import PyPDF2
import re
import os
import csv
import json
import fitz
import unittest

class PdfTest(unittest.TestCase):

    def test_get_pages(self):
        pages = PdfUtil.get_pages(r"C:\Users\zhangyuming\Desktop\公司文件\压力测试\测试个人章位置是否正确\CMMG-YY-202200058中国移动通信集团河南有限公司.pdf")
        print(pages)

    def test_get_text(self):
        filepath = "C:\\Users\\zhangyuming\\Desktop\\公司文件\\压力测试\\测试个人章位置是否正确\\"
        in_filename = "CMMG-YY-202200058中国移动通信集团河南有限公司.pdf"
        out_filename = "CMMG-YY-202200058中国移动通信集团河南有限公司.temp.pdf"
        page = PdfUtil.get_contains_text_page(filepath + in_filename, '甲方：法定代表人或授权代表（签字）')
        print(page)
        PdfUtil.split_single_pdf(filepath + in_filename, page, page + 1, filepath + out_filename)

    def test_page_to_png(self):
        filepath = "C:\\Users\\zhangyuming\\Desktop\\公司文件\\压力测试\\测试个人章位置是否正确\\"
        in_filename = "CMMG-YY-202200058中国移动通信集团河南有限公司.pdf"
        out_filename = "a.png"
        page = PdfUtil.get_contains_text_page(filepath + in_filename, '甲方：法定代表人或授权代表（签字）')
        PdfUtil.save_pdf_page_to_png(filepath + in_filename, page, filepath + out_filename)


class PdfUtil:

    def fun1(self):
        pass

    """
    获取pdf的总行数
    """
    def get_pages(filename):
        pages = 0
        with open(filename, 'rb', ) as fb:
            reader = PyPDF2.PdfFileReader(fb)
            pages = len(reader.pages)
        return pages

    """
    获取pdf文件中某个文本的所在页数
    """
    def get_contains_text_page(filename, text):
        pages = PdfUtil.get_pages(filename)
        with pdfplumber.open(filename) as pdf:
            for i in range(pages):
                # print(pdf.pages[i].extract_text())
                find = re.findall(text, pdf.pages[i].extract_text())
                if find:
                    # return find[0].strip()
                    return i
    """
    将read_file_name对应的pdf的start_page~end_page存储为新的pdf存储到pdf_file
    """
    def split_single_pdf(read_file_name, start_page, end_page, pdf_file):
        # 1. 获取原始pdf文件
        fp_read_file = open(read_file_name, 'rb')
        # 2. 将要分割的PDF内容格式化
        pdf_input = PyPDF2.PdfFileReader(fp_read_file)
        # 3. 实例一个 PDF文件编写器
        pdf_output = PyPDF2.PdfFileWriter()
        # 4. 把3到4页放到PDF文件编写器
        for i in range(start_page, end_page):
            pdf_output.addPage(pdf_input.getPage(i))
        # 5. PDF文件输出
        with open(pdf_file, 'wb') as pdf_out:
            pdf_output.write(pdf_out)
        print(f'{read_file_name}分割{start_page}页-{end_page}页完成，保存为{pdf_file}!')

    """
    将指定的pdf_filename的指定pageNumber页转换为png图片保存到save_png_path路径中
    """
    def save_pdf_page_to_png(pdf_filename, pageNumber, save_png_path):
        doc = fitz.open(pdf_filename)
        page = doc.load_page(pageNumber)  # PDF页数
        pix = page.get_pixmap()
        #pix.writePNG(save_png_path)
        pix.save(save_png_path)

if __name__ == '__main__':
    pass